import { ref, onUnmounted } from 'vue'
import { ElMessage } from 'element-plus'
import type { WebSocketMessage, ProgressUpdate } from '@/types'

export function useWebSocket() {
  const ws = ref<WebSocket | null>(null)
  const isConnected = ref(false)
  const reconnectAttempts = ref(0)
  const maxReconnectAttempts = 5
  const reconnectInterval = 3000

  let reconnectTimer: NodeJS.Timeout | null = null
  let heartbeatTimer: NodeJS.Timeout | null = null
  let messageHandlers: Map<string, ((data: any) => void)[]> = new Map()

  const connect = (taskId: string) => {
    if (ws.value && ws.value.readyState === WebSocket.OPEN) {
      return
    }

    try {
      const wsUrl = `${import.meta.env.VITE_WS_URL || 'ws://localhost:8000'}/api/v1/progress/ws/${taskId}`
      ws.value = new WebSocket(wsUrl)

      ws.value.onopen = () => {
        console.log('WebSocket connected')
        isConnected.value = true
        reconnectAttempts.value = 0

        // 开始心跳
        startHeartbeat()

        // 触发连接成功事件
        emit('connected', { taskId })
      }

      ws.value.onmessage = (event) => {
        try {
          const message: WebSocketMessage = JSON.parse(event.data)
          handleMessage(message)
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error)
        }
      }

      ws.value.onclose = (event) => {
        console.log('WebSocket disconnected:', event.code, event.reason)
        isConnected.value = false
        stopHeartbeat()

        // 触发断开连接事件
        emit('disconnected', { code: event.code, reason: event.reason })

        // 自动重连
        if (event.code !== 1000 && reconnectAttempts.value < maxReconnectAttempts) {
          scheduleReconnect(taskId)
        }
      }

      ws.value.onerror = (error) => {
        console.error('WebSocket error:', error)
        ElMessage.error('WebSocket连接出错')
      }

    } catch (error) {
      console.error('Failed to create WebSocket connection:', error)
      ElMessage.error('无法建立WebSocket连接')
    }
  }

  const disconnect = () => {
    if (reconnectTimer) {
      clearTimeout(reconnectTimer)
      reconnectTimer = null
    }

    stopHeartbeat()

    if (ws.value) {
      ws.value.close(1000, 'Client disconnect')
      ws.value = null
    }

    isConnected.value = false
    reconnectAttempts.value = 0
  }

  const send = (data: any) => {
    if (ws.value && ws.value.readyState === WebSocket.OPEN) {
      ws.value.send(JSON.stringify(data))
    } else {
      console.warn('WebSocket is not connected')
    }
  }

  const on = (event: string, handler: (data: any) => void) => {
    if (!messageHandlers.has(event)) {
      messageHandlers.set(event, [])
    }
    messageHandlers.get(event)!.push(handler)
  }

  const off = (event: string, handler?: (data: any) => void) => {
    if (messageHandlers.has(event)) {
      const handlers = messageHandlers.get(event)!
      if (handler) {
        const index = handlers.indexOf(handler)
        if (index > -1) {
          handlers.splice(index, 1)
        }
      } else {
        handlers.length = 0
      }
    }
  }

  const emit = (event: string, data: any) => {
    if (messageHandlers.has(event)) {
      messageHandlers.get(event)!.forEach(handler => {
        try {
          handler(data)
        } catch (error) {
          console.error('Error in WebSocket event handler:', error)
        }
      })
    }
  }

  const handleMessage = (message: WebSocketMessage) => {
    console.log('WebSocket message received:', message)

    // 触发通用消息事件
    emit('message', message)

    // 触发特定类型事件
    emit(message.type, message.data)

    // 处理特定消息类型
    switch (message.type) {
      case 'connected':
        console.log('WebSocket connection confirmed')
        break

      case 'progress_update':
        handleProgressUpdate(message.data)
        break

      case 'pong':
        // 心跳响应
        break

      case 'error':
        ElMessage.error(message.data.message || '发生错误')
        break
    }
  }

  const handleProgressUpdate = (data: ProgressUpdate) => {
    // 处理进度更新
    emit('progress', data)

    // 如果任务完成，显示通知
    if (data.status === 'completed') {
      ElMessage.success('视频处理完成！')
    } else if (data.status === 'failed') {
      ElMessage.error(`视频处理失败: ${data.error_message || '未知错误'}`)
    }
  }

  const startHeartbeat = () => {
    stopHeartbeat()
    heartbeatTimer = setInterval(() => {
      send({ type: 'ping', timestamp: Date.now() })
    }, 30000) // 每30秒发送一次心跳
  }

  const stopHeartbeat = () => {
    if (heartbeatTimer) {
      clearInterval(heartbeatTimer)
      heartbeatTimer = null
    }
  }

  const scheduleReconnect = (taskId: string) => {
    if (reconnectTimer) {
      return
    }

    reconnectAttempts.value++
    const delay = reconnectInterval * Math.pow(2, reconnectAttempts.value - 1)

    console.log(`Scheduling reconnect attempt ${reconnectAttempts.value} in ${delay}ms`)

    reconnectTimer = setTimeout(() => {
      reconnectTimer = null
      console.log(`Attempting to reconnect (${reconnectAttempts.value}/${maxReconnectAttempts})`)
      connect(taskId)
    }, delay)
  }

  // 组件卸载时清理
  onUnmounted(() => {
    disconnect()
  })

  return {
    ws,
    isConnected,
    reconnectAttempts,
    connect,
    disconnect,
    send,
    on,
    off
  }
}